from flask import Flask, render_template, request
import pandas as pd
import numpy as np
import torch
import torch.nn as nn

app = Flask(__name__)


@app.route('/')
def home():
    """首页"""
    return render_template('home.html')


@app.route('/a')
def age_distribution_bar():
    """年龄分布图表"""
    with open("./templates/a.html", encoding="utf8", mode="r") as f:
        plot = "".join(f.readlines())
    return render_template('index.html', rl=plot)


@app.route('/b')
def gender_pie():
    """性别患病图表"""
    with open("./templates/b.html", encoding="utf8", mode="r") as f:
        plot = "".join(f.readlines())
    return render_template('index.html', rl=plot)


@app.route('/c')
def cardio_rate_by_age_group():
    """年龄组心脏病患病率图表"""
    with open("./templates/c.html", encoding="utf8", mode="r") as f:
        plot = "".join(f.readlines())
    return render_template('index.html', rl=plot)


@app.route('/d')
def health_indicators_analysis():
    """健康状况指标图表"""
    with open("./templates/d.html", encoding="utf8", mode="r") as f:
        plot = "".join(f.readlines())
    return render_template('index.html', rl=plot)


@app.route('/e')
def indicator_radar():
    """主要体征指标图表"""
    with open("./templates/e.html", encoding="utf8", mode="r") as f:
        plot = "".join(f.readlines())
    return render_template('index.html', rl=plot)


@app.route('/f')
def lifestyle_bar():
    """生活习惯与心脏病图表"""
    with open("./templates/f.html", encoding="utf8", mode="r") as f:
        plot = "".join(f.readlines())
    return render_template('index.html', rl=plot)


@app.route('/g')
def feature_correlation_heatmap():
    """特征相关性图表"""
    with open("./templates/g.html", encoding="utf8", mode="r") as f:
        plot = "".join(f.readlines())
    return render_template('index.html', rl=plot)


# 深度学习模型定义
class HeartDiseaseModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(9, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.net(x)


# 加载模型
model = HeartDiseaseModel()
model.load_state_dict(torch.load("./model/heart_model.pth", map_location=torch.device('cpu')))
model.eval()


@app.route('/h')
def predict():
    """预测页面"""
    with open("./templates/h.html", encoding="utf8") as f:
        plot = "".join(f.readlines())
    return render_template('index.html', rl=plot)


@app.route('/predict_result', methods=['POST'])
def predict_result():
    """进行预测并展示结果"""
    try:
        # 获取表单输入
        age = float(request.form['age'])
        sex = float(request.form['sex'])
        ap_hi = float(request.form['ap_hi'])
        ap_lo = float(request.form['ap_lo'])
        cholesterol = float(request.form['cholesterol'])
        smoke = float(request.form['smoke'])
        alco = float(request.form['alco'])
        active = float(request.form['active'])
        height = float(request.form['height'])
        weight = float(request.form['weight'])

        # 计算 BMI 并构建模型输入
        bmi = weight / ((height / 100) ** 2)
        input_tensor = torch.tensor([[age, sex, ap_hi, ap_lo, cholesterol, smoke, alco, active, bmi]],
                                    dtype=torch.float32)

        # 模型预测
        with torch.no_grad():
            prob = model(input_tensor).item()
        result = round(prob * 100, 2)

        return render_template("predict_result.html", prob=result)

    except Exception as e:
        return f"<h3>预测失败：</h3><p>{e}</p><a href='/h'>返回预测</a>"


if __name__ == '__main__':
    app.run(debug=True, port=5001)
